from fabric import Connection, Group
from fabric.exceptions import GroupException
from paramiko import RSAKey
from paramiko.ssh_exception import PasswordRequiredException, SSHException
from os.path import join
import json
import time
from datetime import datetime
from multiprocessing import Pool
from glob import glob
from re import findall, search
from statistics import mean

from benchmark.commands import CommandMaker
from benchmark.config import Committee, NodeParameters, BenchParameters, ConfigError
from benchmark.logs import ParseError, LogParser
from benchmark.utils import BenchError, Print, PathMaker
from honeybadger.instance import InstanceManager

class Bench:
    def __init__(self, ctx):
        self.ctx = ctx
        self.manager = InstanceManager.make()
        self.hosts = self.manager.hosts(flat=True)
        
    def install(self):
        """Install dependencies on all machines"""
        Print.heading('Installing dependencies on all machines')
        
        if not self.hosts:
            raise BenchError('No hosts available')
            
        # Create connection group
        connections = []
        for host in self.hosts:
            try:
                conn = Connection(
                    host=host,
                    user='ubuntu',
                    connect_kwargs={
                        'key_filename': self.manager.settings.key_path,
                        'timeout': 30
                    }
                )
                connections.append(conn)
            except Exception as e:
                Print.error(f'Failed to connect to {host}: {e}')
                continue
        
        if not connections:
            raise BenchError('No valid connections available')
            
        group = Group(*connections)
        
        try:
            # Update system and install dependencies
            Print.info('Updating system packages...')
            group.run('sudo apt update && sudo apt upgrade -y', hide=True)
            
            Print.info('Installing Go...')
            group.run('wget https://go.dev/dl/go1.21.0.linux-amd64.tar.gz', hide=True)
            group.run('sudo tar -C /usr/local -xzf go1.21.0.linux-amd64.tar.gz', hide=True)
            group.run('echo "export PATH=$PATH:/usr/local/go/bin" >> ~/.bashrc', hide=True)
            group.run('export PATH=$PATH:/usr/local/go/bin', hide=True)
            
            Print.info('Installing tmux...')
            group.run('sudo apt install -y tmux', hide=True)
            
            Print.heading('Installation completed successfully')
            
        except GroupException as e:
            Print.error(f'Installation failed: {e}')
            raise BenchError('Failed to install dependencies', e)
        finally:
            for conn in connections:
                conn.close()
    
    def upload_exec(self):
        """Upload the compiled executable to all machines"""
        Print.heading('Uploading executable to all machines')
        
        if not self.hosts:
            raise BenchError('No hosts available')
            
        # First compile locally
        Print.info('Compiling executable...')
        cmd = CommandMaker.compile().split()
        import subprocess
        subprocess.run(cmd, check=True)
        
        # Upload to each host
        for host in self.hosts:
            try:
                conn = Connection(
                    host=host,
                    user='ubuntu',
                    connect_kwargs={
                        'key_filename': self.manager.settings.key_path,
                        'timeout': 30
                    }
                )
                
                Print.info(f'Uploading to {host}...')
                conn.put('main', 'main')
                conn.run('chmod +x main')
                conn.close()
                
            except Exception as e:
                Print.error(f'Failed to upload to {host}: {e}')
                continue
                
        Print.heading('Upload completed')
    
    def run(self, bench_params, node_params, debug=False):
        """Run benchmark on remote machines"""
        Print.heading('Starting remote benchmark')
        
        try:
            bench_parameters = BenchParameters(bench_params)
            node_parameters = NodeParameters(node_params)
        except ConfigError as e:
            raise BenchError('Invalid parameters', e)
        
        if not self.hosts:
            raise BenchError('No hosts available')
            
        # Use first host as coordinator
        coordinator = self.hosts[0]
        workers = self.hosts[1:]
        
        try:
            # Connect to coordinator
            coord_conn = Connection(
                host=coordinator,
                user='ubuntu',
                connect_kwargs={
                    'key_filename': self.manager.settings.key_path,
                    'timeout': 30
                }
            )
            
            # Cleanup and prepare
            Print.info('Preparing coordinator...')
            coord_conn.run(CommandMaker.cleanup_configs(), hide=True)
            
            # Generate keys
            nodes = bench_parameters.nodes[0]
            Print.info(f'Generating keys for {nodes} nodes...')
            coord_conn.run(CommandMaker.generate_key(path="./", nodes=nodes), hide=True)
            
            # Generate committee
            Print.info('Generating committee...')
            committee_data = self._generate_committee(coord_conn, nodes)
            
            # Generate parameters
            Print.info('Generating parameters...')
            coord_conn.run(f'echo \'{json.dumps(node_parameters.json)}\' > .parameters.json', hide=True)
            
            # Start nodes
            Print.info('Starting nodes...')
            for i in range(nodes):
                host = self.hosts[i % len(self.hosts)]
                if host == coordinator:
                    conn = coord_conn
                else:
                    conn = Connection(
                        host=host,
                        user='ubuntu',
                        connect_kwargs={
                            'key_filename': self.manager.settings.key_path,
                            'timeout': 30
                        }
                    )
                
                # Upload necessary files
                if host != coordinator:
                    conn.put(f'key-{i}.json', f'key-{i}.json')
                    conn.put('.committee.json', '.committee.json')
                    conn.put('.parameters.json', '.parameters.json')
                
                # Start node in tmux
                cmd = CommandMaker.run_node(
                    i,
                    f'key-{i}.json',
                    f'threshold-key-{i}.json',
                    '.committee.json',
                    f'db-{i}',
                    '.parameters.json',
                    "remote",
                    bench_parameters.log_level
                )
                
                conn.run(f'tmux new-session -d -s node-{i} "{cmd}"', hide=True)
                
                if host != coordinator:
                    conn.close()
            
            # Wait for benchmark
            Print.info(f'Running benchmark for {bench_parameters.duration} seconds...')
            time.sleep(bench_parameters.duration)
            
            # Stop all nodes
            Print.info('Stopping nodes...')
            for i in range(nodes):
                host = self.hosts[i % len(self.hosts)]
                if host == coordinator:
                    conn = coord_conn
                else:
                    conn = Connection(
                        host=host,
                        user='ubuntu',
                        connect_kwargs={
                            'key_filename': self.manager.settings.key_path,
                            'timeout': 30
                        }
                    )
                
                conn.run(f'tmux kill-session -t node-{i}', hide=True, warn=True)
                
                if host != coordinator:
                    conn.close()
            
            # Download logs
            Print.info('Downloading logs...')
            self._download_logs(coord_conn, "remote")
            
            # Parse results
            Print.info('Parsing logs...')
            try:
                parser = LogParser.process('./logs/remote', 
                                        faults=node_parameters.faults,
                                        protocol=bench_parameters.protocol,
                                        ddos=node_parameters.ddos)
                return parser
            except ParseError as e:
                raise BenchError('Failed to parse logs', e)
                
        except Exception as e:
            raise BenchError('Benchmark failed', e)
        finally:
            if 'coord_conn' in locals():
                coord_conn.close()
    
    def _generate_committee(self, conn, nodes):
        """Generate committee configuration"""
        # Read public keys
        keys = []
        for i in range(nodes):
            result = conn.run(f'cat key-{i}.json', hide=True)
            data = json.loads(result.stdout)
            keys.append(data['public'])
        
        # Generate addresses
        addresses = []
        for i in range(nodes):
            host = self.hosts[i % len(self.hosts)]
            port = 8000 + i
            addresses.append(f'{host}:{port}')
        
        # Create committee
        committee = Committee(keys, list(range(nodes)), addresses)
        committee.print('.committee.json')
        
        return committee
    
    def _download_logs(self, conn, timestamp):
        """Download logs from coordinator"""
        import os
        os.makedirs(f'./logs/{timestamp}', exist_ok=True)
        
        # Download log files
        result = conn.run('ls logs/', hide=True)
        if result.return_code == 0:
            log_files = result.stdout.strip().split('\n')
            for log_file in log_files:
                if log_file.strip():
                    conn.get(f'logs/{log_file}', f'./logs/{timestamp}/{log_file}')
    
    def kill(self):
        """Kill all running processes on all machines"""
        Print.heading('Killing all processes on all machines')
        
        for host in self.hosts:
            try:
                conn = Connection(
                    host=host,
                    user='ubuntu',
                    connect_kwargs={
                        'key_filename': self.manager.settings.key_path,
                        'timeout': 30
                    }
                )
                
                conn.run(CommandMaker.kill(), hide=True, warn=True)
                conn.close()
                
            except Exception as e:
                Print.error(f'Failed to kill processes on {host}: {e}')
                continue
                
        Print.heading('All processes killed')
    
    def download(self, node_instance, timestamp):
        """Download logs from specific node instance"""
        Print.heading(f'Downloading logs from node instance {node_instance}')
        
        if not self.hosts:
            raise BenchError('No hosts available')
            
        host = self.hosts[node_instance % len(self.hosts)]
        
        try:
            conn = Connection(
                host=host,
                user='ubuntu',
                connect_kwargs={
                    'key_filename': self.manager.settings.key_path,
                    'timeout': 30
                }
            )
            
            # Create local directory
            import os
            os.makedirs(f'./logs/{timestamp}', exist_ok=True)
            
            # Download logs
            result = conn.run('ls logs/', hide=True)
            if result.return_code == 0:
                log_files = result.stdout.strip().split('\n')
                for log_file in log_files:
                    if log_file.strip():
                        conn.get(f'logs/{log_file}', f'./logs/{timestamp}/{log_file}')
            
            conn.close()
            return f'Logs downloaded from {host}'
            
        except Exception as e:
            raise BenchError(f'Failed to download logs from {host}', e)

